from __future__ import annotations

from typing import Dict, Any, List
import json
import re
import os

from ..utils.ai_utils import AIClient


SYSTEM_PROMPT = (
    "你是一个资深数据库助理。用户用自然语言提问，你需要基于提供的'可用表与列'生成一条可直接执行的 SQL。"
    "生成规则："
    "1) 只使用提供的表和字段，不要编造字段名或表名；优先 SELECT 查询；标识符引用请根据数据库类型选择："
    "   - MySQL 使用反引号`name`；PostgreSQL 使用双引号\"name\"；SQLite 可直接使用裸标识符。并添加合理 LIMIT。"
    "2) 根据字段名语义自行判断应聚合的指标（如金额、数量等），但必须从提供的字段中选择；避免选择与语义不符的字段。"
    "3) 需要联表时，依据字段名相似度推断主外键（如 *_id、*编码、*编号、id 等）进行 JOIN；只在提供字段中存在时联表。"
    "4) 尽量以事实表（如订单/明细/销售流水）聚合，并可关联维度表（产品/门店等）取名称；注意 WHERE/时间范围与 GROUP BY 的正确性。"
    "5) 在生成WHERE条件前，必须仔细分析样本数据中每列的实际值，理解列的真实含义。例如："
    "   - 如果查询包含人名（如'周健'），先检查样本数据中哪些列包含人名，哪些列包含公司名/部门名"
    "   - 如果查询包含地名（如'北京'），先检查样本数据中哪些列包含地名，哪些列包含其他信息"
    "   - 选择WHERE条件时，优先使用语义匹配的列，避免使用语义不匹配的列"
    "输出：只返回 SQL 一行，不要解释，也不要带```代码块。"
)


class MultiAgentOrchestrator:
    def __init__(self, model: str | None = None):
        # 默认对齐通义千问 DashScope 兼容接口
        self.model = model or (os.getenv("OPENAI_MODEL") or "qwen3-235b-a22b")

    def extract_intent(self, query: str) -> Dict[str, Any]:
        """让模型先抽取意图与可能的过滤值，返回结构化 JSON。容错解析，失败则返回空。"""
        prompt = (
            "请把用户的自然语言查询解析为一个紧凑 JSON，不要解释。"
            "字段：metrics_keywords(数组)、dimensions_keywords(数组)、filter_values(数组)、time_range(字符串或空)。"
            "只返回 JSON。"
        )
        content = AIClient.chat(
            model=self.model,
            messages=[
                {"role": "system", "content": prompt},
                {"role": "user", "content": query},
            ],
            temperature=0.1,
            max_tokens=200,
        )
        try:
            # 提取首个代码块/或直接 JSON
            m = re.search(r"```[a-zA-Z]*\n([\s\S]*?)\n```", content)
            txt = m.group(1).strip() if m else content.strip()
            data = json.loads(txt)
            if isinstance(data, dict):
                return data
        except Exception:
            pass
        return {"metrics_keywords": [], "dimensions_keywords": [], "filter_values": [], "time_range": ""}

    def generate_sql(self, query: str, db_type: str = "mysql", available_tables: list[str] | None = None, table_schemas: dict[str, list[str]] | None = None, table_samples: dict[str, list[dict]] | None = None) -> str:
        # 组织可用表/列的提示
        tables_text = "无"
        if table_schemas:
            parts = []
            for t, cols in table_schemas.items():
                if (db_type or "").lower() == "postgresql":
                    preview = ", ".join([f'"{c}"' for c in cols[:24]])
                    parts.append(f"- \"{t}\": {preview}")
                else:
                    preview = ", ".join([f"`{c}`" for c in cols[:24]])
                    parts.append(f"- `{t}`: {preview}")
            tables_text = "\n".join(parts) if parts else tables_text
        elif available_tables:
            if (db_type or "").lower() == "postgresql":
                tables_text = ", ".join([f'"{t}"' for t in available_tables])
            else:
                tables_text = ", ".join([f"`{t}`" for t in available_tables])

        # 拼接轻量样本，帮助模型理解“北京”等值可能出现在地址等字段
        samples_text = ""
        candidate_filter_cols: List[str] = []
        if table_samples:
            parts = []
            # 先抽取意图，得到更明确的过滤值
            intent = self.extract_intent(query or "")
            intent_values = []
            if isinstance(intent, dict):
                if isinstance(intent.get("filter_values"), list):
                    intent_values.extend([str(x) for x in intent["filter_values"] if x is not None])
            # 从查询中提取 2-6 长度的中文片段作为补充（如“北京”）
            query_cn_tokens = list(set(re.findall(r"[\u4e00-\u9fa5]{2,6}", query or "")))
            if intent_values:
                query_cn_tokens = list(set(query_cn_tokens + intent_values))
            for t, rows in table_samples.items():
                if not rows:
                    continue
                # 限制样本数据长度，避免数据过大
                sample_json = json.dumps(rows, ensure_ascii=False, default=str)
                if len(sample_json) > 3000:  # 如果超过3000字符，截取前几行
                    sample_json = json.dumps(rows[:5], ensure_ascii=False, default=str)
                    if len(sample_json) > 3000:
                        sample_json = sample_json[:3000] + "..."
                parts.append(f"- `{t}` 样本: {sample_json}")
                # 动态识别候选筛选列：仅基于样本值与查询文本，不依赖写死关键词
                matched_any = False
                text_like_cols = set()
                for row in rows:
                    for col, val in row.items():
                        sval = "" if val is None else str(val)
                        # 识别文本型列（包含中文或字母，且含空格/分隔符/长度较长），用于无匹配时的兜底
                        if re.search(r"[\u4e00-\u9fa5A-Za-z]", sval) and len(sval) >= 2:
                            if (db_type or "").lower() == "postgresql":
                                text_like_cols.add(f'"{t}"."{col}"')
                            else:
                                text_like_cols.add(f"`{t}`.`{col}`")
                        # 如果样本值包含查询中的中文片段，则认为该列是候选筛选列
                        if any(tok and tok in sval for tok in query_cn_tokens):
                            if (db_type or "").lower() == "postgresql":
                                candidate_filter_cols.append(f'"{t}"."{col}"')
                            else:
                                candidate_filter_cols.append(f"`{t}`.`{col}`")
                            matched_any = True
                # 若该表未匹配到查询片段，则选取少量文本型列作为兜底候选
                if not matched_any and text_like_cols:
                    for c in list(text_like_cols):
                        candidate_filter_cols.append(c)
            samples_text = "\n".join(parts)
        # 去重并限制长度
        candidate_filter_cols = list(dict.fromkeys(candidate_filter_cols))
        content = AIClient.chat(
            model=self.model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT + "\n\n重要：请仔细分析样本数据中每列的实际值，理解列的含义。例如：如果某列存放的是公司名称而不是人名，就不要用它来筛选人名。"},
                {"role": "user", "content": f"数据库类型: {db_type}\n可用表与列:\n{tables_text}\n{('样本数据:\n' + samples_text + '\n') if samples_text else ''}{('候选筛选字段(可用于 WHERE，优先使用 LIKE): ' + ', '.join(candidate_filter_cols) + '\n') if candidate_filter_cols else ''}\n\n在生成SQL前，请先分析样本数据：\n1. 仔细查看每列的实际值\n2. 理解每列的含义（如：人名、公司名、地名等）\n3. 选择语义匹配的列进行WHERE筛选\n\n需求: {query}"},
            ],
            temperature=0.1,
            max_tokens=800,  # 增加token数量，让大模型有更多空间处理
            timeout=1800,    # 增加超时时间到30分钟
        )
        # 提取首个代码块中的 SQL 或整段文本
        match = re.search(r"```[a-zA-Z]*\n([\s\S]*?)\n```", content)
        sql = match.group(1).strip() if match else content.strip()
        return sql

    def generate_sql_with_retry(self, query: str, db_type: str = "mysql", available_tables: list[str] | None = None, table_schemas: dict[str, list[str]] | None = None, table_samples: dict[str, list[dict]] | None = None, failed_sql: str = "", retry_count: int = 0) -> str:
        """重新生成SQL，让大模型基于样本数据自己分析并选择更好的列"""
        # 构建改进的提示
        tables_text = "无"
        if table_schemas:
            parts = []
            for t, cols in table_schemas.items():
                preview = ", ".join([f"`{c}`" for c in cols[:24]])
                parts.append(f"- `{t}`: {preview}")
            tables_text = "\n".join(parts) if parts else tables_text
        elif available_tables:
            tables_text = ", ".join([f"`{t}`" for t in available_tables])

        # 拼接样本数据
        samples_text = ""
        candidate_filter_cols: List[str] = []
        if table_samples:
            parts = []
            # 先抽取意图，得到更明确的过滤值
            intent = self.extract_intent(query or "")
            intent_values = []
            if isinstance(intent, dict):
                if isinstance(intent.get("filter_values"), list):
                    intent_values.extend([str(x) for x in intent["filter_values"] if x is not None])
            # 从查询中提取 2-6 长度的中文片段作为补充（如"北京"）
            query_cn_tokens = list(set(re.findall(r"[\u4e00-\u9fa5]{2,6}", query or "")))
            if intent_values:
                query_cn_tokens = list(set(query_cn_tokens + intent_values))
            
            for t, rows in table_samples.items():
                if not rows:
                    continue
                # 限制样本数据长度，避免数据过大
                sample_json = json.dumps(rows, ensure_ascii=False, default=str)
                if len(sample_json) > 3000:  # 如果超过3000字符，截取前几行
                    sample_json = json.dumps(rows[:5], ensure_ascii=False, default=str)
                    if len(sample_json) > 3000:
                        sample_json = sample_json[:3000] + "..."
                parts.append(f"- `{t}` 样本: {sample_json}")
                # 动态识别候选筛选列
                matched_any = False
                text_like_cols = set()
                for row in rows:
                    for col, val in row.items():
                        sval = "" if val is None else str(val)
                        # 识别文本型列（包含中文或字母，且含空格/分隔符/长度较长），用于无匹配时的兜底
                        if re.search(r"[\u4e00-\u9fa5A-Za-z]", sval) and len(sval) >= 2:
                            text_like_cols.add(f"`{t}`.`{col}`")
                        # 如果样本值包含查询中的中文片段，则认为该列是候选筛选列
                        if any(tok and tok in sval for tok in query_cn_tokens):
                            candidate_filter_cols.append(f"`{t}`.`{col}`")
                            matched_any = True
                # 若该表未匹配到查询片段，则选取少量文本型列作为兜底候选
                if not matched_any and text_like_cols:
                    for c in list(text_like_cols):
                        candidate_filter_cols.append(c)
            samples_text = "\n".join(parts)
        
        # 去重
        candidate_filter_cols = list(dict.fromkeys(candidate_filter_cols))
        
        # 构建重试提示，让大模型自己分析样本数据
        retry_prompt = f"数据库类型: {db_type}\n可用表与列:\n{tables_text}\n{('样本数据:\n' + samples_text + '\n') if samples_text else ''}{('候选筛选字段(可用于 WHERE，优先使用 LIKE): ' + ', '.join(candidate_filter_cols) + '\n') if candidate_filter_cols else ''}\n需求: {query}"
        
        if failed_sql:
            retry_prompt += f"\n\n注意：上次SQL执行失败或结果为空，请仔细分析样本数据，选择包含有效值的列进行筛选。"

        content = AIClient.chat(
            model=self.model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT + "\n\n重要：请仔细分析样本数据中每列的实际值，理解列的含义。例如：如果某列存放的是公司名称而不是人名，就不要用它来筛选人名。"},
                {"role": "user", "content": retry_prompt},
            ],
            temperature=0.1,
            max_tokens=800,  # 增加token数量
            timeout=1800,    # 增加超时时间到30分钟
        )
        
        # 提取首个代码块中的 SQL 或整段文本
        match = re.search(r"```[a-zA-Z]*\n([\s\S]*?)\n```", content)
        sql = match.group(1).strip() if match else content.strip()
        return sql


